import time
import numpy as np
import sys
import os
from config import port
# IKPy is just an example real IK library.
# You need to install it (e.g., pip install ikpy)
# and set up your chain parameters (DH or URDF).
from ikpy.chain import Chain #type: ignore
from ikpy.link import URDFLink #type: ignore
from DepthColorProcessor import DepthColorProcessor
import numpy as np
import pyrealsense2 as rs#type: ignore
from scipy.spatial.transform import Rotation as R#type: ignore
from bin.camtobase import cam_to_base_link_pose

# Insert the path containing serialControl.py
taskagent_path = os.path.dirname(__file__)  # or a known absolute path
print(taskagent_path)
if taskagent_path not in sys.path:
    sys.path.insert(0, taskagent_path)
from Serial_Control import SerialControl  # Adjust path if needed

class armMotion:
    """
    A class for controlling a 5-DOF manipulator without ROS dependencies.
    References logic found in xuanwu_arm_control.cpp:
      - is_planning_success -> tries an IK solution

    Uses ikpy for a 'real' IK calculation.
    """

    def __init__(self,ser:SerialControl,dcp:DepthColorProcessor):
        """
        Initialize poses, serial connection, and an IK chain.
        """
        self.dcp = dcp
        self.ser = ser
        self.init_joints = [0.09189158511750145, -1.1898782175471343, 1.7011724219188729, 1.6069246423111792, 0.0]
        self.last_plan_joints = self.init_joints
        # self.place_joints = [-0.04948, 0.53957, -1.75301, -1.05793, 0.0]
        self.pack_pose = (0.15,0.0,0.15)
        self.straight_pose = (0.0,0.0,0.5)
        self.openCommand = "#026P1200T0700!"
        self.pickCommand = "#026P1900T0700!"
        self.joint_rads = [0.0, 0.0, 0.0, 0.0, 0.0]
        self.chain = Chain.from_urdf_file("arm.urdf")
        # self.moveToInit()
        self.firstCall=True
        print("[armMotion] Initialization complete.")
    def plan(self, pos: tuple[float, float, float]):
        joint_rads = []
        print("[is_planning_success] Attempting to solve IK with ikpy...")
        try:
            # Pass the mask to inverse_kinematics:
            ik_solution = self.chain.inverse_kinematics(
                pos,
                [0,0,0,1],
            )
            ik_solution = ik_solution[1:6]

            print("[plan] Found IK solution:")
            for i, angle in enumerate(ik_solution):
                joint_rads.append(angle)
                print(f"  joint_rads[{i}] = {angle:.5f}")
            return joint_rads
        except Exception as e:
            print(f"[plan] IK failed: {e}")
            return []
    
    def moveToPose(self, 
                   pose: tuple[float, float, float]):
        """
        If plan_joints is provided, skip the IK planning step.
        Otherwise, compute it from the given pose.
        """
        # If no joint plan is provided, compute it now:
        plan_joints = self.plan(pose)
        if len(plan_joints) == 0:
            return False
        # Example offset and forced constraints
        plan_joints[0] += 0.11781
        plan_joints[4] = 0.0
        self.last_plan_joints = plan_joints
        servoRads=[]
        # Send the joints to your servo controller
        for i in range(len(plan_joints)):
            servoRads.append([i+21,plan_joints[i]])
        self.ser.servoRads(servoRads, True)
        return True
    def moveToUxy(self, 
                     pos_2d: tuple[int, int]):
        """
        If plan_joints is provided, skip the IK planning step after computing (x,y,z).
        Otherwise, plan internally.
        
        This also uses a local matrix (T_base_to_end_2D) rather than self.T_base_to_end.
        """
        depth_frame = self.dcp.depth_bin_to_png("depth_data.bin", "output_depth.png")
        # Example of a local transform if you have or compute it at runtime
        T_base_to_end_2D = np.eye(4)  # or some other transform you compute locally
        x, y, z = cam_to_base_link_pose(
            pos_2d[0],
            pos_2d[1],
            depth_frame[pos_2d[0]][pos_2d[1]],
            T_base_to_end_2D
        )
        # Only plan the IK if no precomputed plan was passed in:
        return self.moveToPose(pose=(x, y, z))
    def openGripper(self):
        self.ser.sendStr(self.openCommand)       # Open gripper
    def closeGripper(self):
        self.ser.sendStr(self.pickCommand)       # Close gripper
    def moveToInit(self):
        self.ser.servoRads(self.init_joints)# Return to init
    def moveToPack(self):
        self.moveToPose(self.pack_pose)
    def moveToStraight(self):
        self.moveToPose(self.straight_pose)
    def pickAndPlace(self,uxuy,placeIt=True,placePos=None):
        sleepTime=1
        self.openGripper()
        time.sleep(sleepTime)
        self.moveToUxy(uxuy)
        time.sleep(sleepTime)
        self.closeGripper()
        time.sleep(sleepTime)
        if placeIt:
            if placePos is not None:
                self.moveToPose(placePos)
            else:
                self.moveToPack()
            time.sleep(sleepTime)
            self.openGripper()
            time.sleep(sleepTime)
            self.moveToInit()
    def toPos(self,x,y,z):
        self.moveToPose((x,y,z))
        time.sleep(2)
    def catch(self):
        self.closeGripper()
        time.sleep(2)
    def release(self):
        self.openGripper()
        time.sleep(2)
    def toDir(self,yaw,pitch):
        if self.firstCall:
            self.moveToStraight()
            self.firstCall=False
        yaw=yaw*np.pi/180
        pitch=pitch*np.pi/180
        self.ser.servoRads([(21,yaw),(24,pitch)])
        time.sleep(2)
    def grabThings(self):
        self.moveToInit()
        time.sleep(2)
        self.openGripper()
        time.sleep(2)
        # self.moveToPose((-0.3,0,0.06))
        self.moveToPose((-0.2,0.1,-0.05))
        time.sleep(2)
        self.closeGripper()
        time.sleep(1)
        self.moveToStraight()
    def forwardK(self) -> tuple[tuple[float,float,float], list[float]]:
        """
        Compute forward kinematics from the last planned joint angles.

        Returns:
            (trans, quat_xyzw):
                trans: (x, y, z) as a tuple, representing end-effector translation in the base frame
                quat_xyzw: [x, y, z, w] as a list, representing end-effector orientation in the base frame
        """
        if not self.last_plan_joints:
            print("[forwardK] No last_plan_joints available. Returning None.")
            return None, None#type: ignore
        
        # 'ikpy' expects a full joint set: typically the chain might have 6 or 7 moving joints.
        # If your chain has an extra joint for the base or end, prepend/append 0.0 or use
        # your specific number of joints to match IK solutions. This example uses:
        joint_vector = [0.0] + self.last_plan_joints + [0.0]

        # Compute the forward kinematics transform (4×4 matrix)
        T = self.chain.forward_kinematics(joint_vector)

        # Extract translation and rotation
        trans = (T[0, 3], T[1, 3], T[2, 3])
        rot_mat_3x3 = T[:3, :3]
        r_obj = R.from_matrix(rot_mat_3x3)
        quat_xyzw = r_obj.as_quat().tolist()  # [x, y, z, w]

        # Print / debug if you want to see the results
        print("[forwardK] FK end-effector pose:")
        print(f"  trans = {trans}")
        print(f"  quat_xyzw = {quat_xyzw}")
        return trans, quat_xyzw
if __name__ == "__main__":  
    controller = armMotion(SerialControl(port, 115200, 1000),DepthColorProcessor(interactive_plots=False))
    while True:
        # ux,uy=input("Enter ux,uy: ").split(",")
        # controller.pickAndPlace((int(ux),int(uy)))
        # controller.grabThings()
        try:
            command=input("Enter x,y,z or g or r: ")
            if command=="g":
                controller.catch()
            elif command=="r":
                controller.release()
            else:
                x,y,z=command.split(",")
                controller.moveToPose((float(x),float(y),float(z)))
        except Exception as e:
            print(e)
